use std::collections::HashMap;
use std::rc::Rc;
use std::fmt::Display;
use std::fmt::Formatter;
use syn::File;
use syn::Item;

// sympol table manager
struct Table {
    sympols: HashMap<String, Reg>,
}

impl Table {
    fn new() -> Self {
        Self {
            sympols: HashMap::new(),
        }
    }

    fn sympol_add(&mut self, ident: String, reg: Reg) {
        self.sympols.insert(ident, reg);
    }

    fn sympol_get(&self, ident: &str) -> Option<Reg> {
        self.sympols.get(ident).copied()
    }
}

struct TableManager {
    tables: Vec<Table>,
}

impl TableManager {
    fn new() -> Self {
        Self {
            tables: Vec::new(),
        }
    }

    fn scope_push(&mut self) {
        self.tables.push(Table::new());
    }

    fn scope_pop(&mut self) {
        self.tables.pop().unwrap();
    }

    fn sympol_add(&mut self, ident: String, reg: Reg) {
        self.tables.last_mut().unwrap().sympol_add(ident, reg);
    }

    fn sympol_get(&self, ident: &str) -> Option<Reg> {
        for table in &self.tables {
            if let Some(reg) = table.sympol_get(ident) {
                return Some(reg)
            }
        }
        None
    }
}

#[derive(Clone, Copy)]
struct Reg(usize);

impl Display for Reg {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(f, "_{}", self.0)
    }
}

enum Opt {
    Reg(Reg),
    Imm(i64),
}

impl From<Reg> for Opt {
    fn from(value: Reg) -> Self {
        Self::Reg(value)
    }
}

impl From<i64> for Opt {
    fn from(value: i64) -> Self {
        Self::Imm(value)
    }
}

impl Display for Opt {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        match self {
            Opt::Reg(reg) => write!(f, "{}", reg),
            Opt::Imm(imm) => write!(f, "{}", imm),
        }
    }
}

enum InstOption {
    Ass,
    Add,
    Sub,
    Mul,
    Gt,
    Lt,
    Ge,
    Le,
    Eq,
}

struct Inst {
    inst: InstOption,
    rd: Reg,
    rs1: Opt,
    rs2: Opt,
}

impl Display for Inst {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        match self.inst {
            InstOption::Ass => write!(f, "{} = {};", self.rd, self.rs1),
            InstOption::Add => write!(f, "{} = Add({}, {});", self.rd, self.rs1, self.rs2),
            InstOption::Sub => write!(f, "{} = Sub({}, {});", self.rd, self.rs1, self.rs2),
            InstOption::Mul => write!(f, "{} = Mul({}, {});", self.rd, self.rs1, self.rs2),
            InstOption::Gt => write!(f, "{} = Gt({}, {});", self.rd, self.rs1, self.rs2),
            InstOption::Lt => write!(f, "{} = Lt({}, {});", self.rd, self.rs1, self.rs2),
            InstOption::Ge => write!(f, "{} = Ge({}, {});", self.rd, self.rs1, self.rs2),
            InstOption::Le => write!(f, "{} = Le({}, {});", self.rd, self.rs1, self.rs2),
            InstOption::Eq => write!(f, "{} = Eq({}, {});", self.rd, self.rs1, self.rs2),
        }
    }
}

struct SwitchInt {
    rs: Reg,
    true_res: bool,
    true_block: Rc<BasicBlock>,
    false_block: Rc<BasicBlock>,
}

impl Display for SwitchInt {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(f, "SwitchInt({}) [{}: {}, otherwise: {}];", self.rs, if self.true_res {1} else {0}, self.true_block.get_name(), self.false_block.get_name())
    }
}

struct Goto {
    target: Box<BasicBlock>,
}

impl Display for Goto {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(f, "goto -> {}", self.target.get_name())
    }
}

enum BlockEnd {
    SwitchInt(SwitchInt),
    Goto(Goto),
    Return,
}

impl Display for BlockEnd {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        match self {
            BlockEnd::SwitchInt(switch_int) => write!(f, "{}", switch_int),
            BlockEnd::Goto(goto) => write!(f, "{}", goto),
            BlockEnd::Return => write!(f, "return;"),
        }
    }
}

struct BasicBlock {
    id: usize,
    insts: Vec<Inst>,
    block_end: Option<BlockEnd>,
}

impl BasicBlock {
    fn new(id: usize) -> Self {
        Self {
            id,
            insts: Vec::new(),
            block_end: None,
        }
    }

    fn get_name(&self) -> String {
        format!("bb{}", self.id)
    }

    fn push_inst(&mut self, inst: Inst) {
        self.insts.push(inst);
    }

    fn set_end(&mut self, block_end: BlockEnd) {
        self.block_end = Some(block_end);
    }
}

impl Display for BasicBlock {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        writeln!(f, "    bb{}: {{", self.id)?;
        for inst in &self.insts {
            writeln!(f, "        {}", inst)?;
        }
        writeln!(f, "        {}", self.block_end.as_ref().unwrap())?;
        writeln!(f, "    }}")
    }
}

pub struct MIRFunction {
    ident: String,
    args: Vec<(String, Reg)>,
    have_ret: bool,
    basic_blocks: Vec<Rc<BasicBlock>>,
}

impl MIRFunction {
    fn push_block(&mut self, block: Rc<BasicBlock>) {
        self.basic_blocks.push(block);
    }
}

impl Display for MIRFunction {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(f, "fn {}", self.ident)?;
        write!(f, "(")?;
        let mut args_str: String = String::new();
        for (_, reg) in &self.args {
            args_str.push_str(&format!("{}: i32", reg));
            args_str.push_str(", ");
        }
        if !args_str.is_empty() {
            args_str.pop().unwrap();
            args_str.pop().unwrap();
        }
        write!(f, "{}", args_str)?;
        write!(f, ")")?;
        write!(f, " -> ")?;
        if self.have_ret {
            write!(f, "i32")?;
        } else {
            write!(f, "()")?;
        }
        write!(f, " ")?;
        writeln!(f, "{{")?;
        let mut new_line = false;
        // debug
        for (ident, reg) in &self.args {
            writeln!(f, "    debug {} => {};", ident, reg)?;
            new_line = true;
        }
        // block
        for block in &self.basic_blocks {
            if new_line {
                writeln!(f, "")?;
            } else {
                new_line = true;
            }
            write!(f, "{}", block)?;
        }
        writeln!(f, "}}")
    }
}

struct MIRGen {
    func: MIRFunction,
    tm: TableManager,
    reg_id: usize,
    block_id: usize,
}

impl MIRGen {
    fn new(ident: String, args: Vec<String>, have_ret: bool) -> Self {
        // create args' reg
        let mut reg_id = 1;
        let mut func_args: Vec<(String, Reg)> = Vec::new();
        for arg in args {
            let reg: Reg = Reg(reg_id);
            let ident: String = arg;
            func_args.push((ident, reg));
            reg_id += 1;
        }
        // add args to table
        let mut tm = TableManager::new();
        tm.scope_push();
        for arg in &func_args {
            tm.sympol_add(arg.0.clone(), arg.1);
        }

        let func = MIRFunction {
            ident,
            args: func_args,
            have_ret,
            basic_blocks: Vec::new(),
        };
        Self {
            func,
            tm,
            reg_id,
            block_id: 0,
        }
    }

    fn create_reg(&mut self) -> Reg {
        let reg = Reg(self.reg_id);
        self.reg_id += 1;
        reg
    }

    fn create_block(&mut self) -> BasicBlock {
        let block_id = self.block_id;
        self.block_id += 1;
        BasicBlock::new(block_id)
    }

    fn push_block(&mut self, block: Rc<BasicBlock>) {
        self.func.push_block(block);
    }
}

pub fn mir_gen(input: &str) -> MIRFunction {
    let syntex: File = syn::parse_file(input).unwrap();
    let (mut mg, block) = match syntex.items.iter().next().unwrap() {
        Item::Fn(func) => {
            let sig = &func.sig;
            let block = func.block.clone();
            let ident = sig.ident.to_string();
            let inputs = &sig.inputs;
            let mut args: Vec<String> = Vec::new();
            for input in inputs {
                let arg_ident = match input {
                    syn::FnArg::Typed(pat_type) => {
                        match &*pat_type.pat {
                            syn::Pat::Ident(ident) => &ident.ident,
                            _ => unreachable!(),
                        }
                    }
                    _ => unreachable!(),
                };
                args.push(arg_ident.to_string());
            }
            let have_ret = match sig.output {
                syn::ReturnType::Default => false,
                syn::ReturnType::Type(_, _) => true,
            };
            (MIRGen::new(ident, args, have_ret), block)
        }
        _ => unreachable!(),
    };
    let mut this_block = mg.create_block();
    visit_block(*block, &mut this_block, &mut mg);
    mg.push_block(Rc::new(this_block));
    mg.func.basic_blocks.sort_by_key(|b| b.id);
    mg.func
}

fn visit_block(block: syn::Block, this_block: &mut BasicBlock, mg: &mut MIRGen) {
    mg.tm.scope_push();
    for stmt in block.stmts {
        match stmt {
            syn::Stmt::Local(local) => {
                // get ident
                let ident = match local.pat {
                    syn::Pat::Ident(ident) => ident.ident,
                    _ => unreachable!(),
                };
                // get init_val
                let reg = match visit_expr(*local.init.unwrap().expr, this_block, mg) {
                    Opt::Reg(reg) => reg,
                    Opt::Imm(imm) => {
                        let reg = mg.create_reg();
                        let inst = Inst {
                            inst: InstOption::Ass,
                            rd: reg,
                            rs1: imm.into(),
                            rs2: imm.into(),
                        };
                        this_block.push_inst(inst);
                        reg
                    }
                };

                mg.tm.sympol_add(ident.to_string(), reg);
            }
            syn::Stmt::Expr(expr, _) => {
                match expr {
                    syn::Expr::If(if_expr) => {
                        let mut true_block = mg.create_block();
                        let mut false_block = mg.create_block();
                        let reg = visit_cond_expr(*if_expr.cond, this_block, mg);
                        visit_block(if_expr.then_branch, &mut true_block, mg);
                        let else_block = match *if_expr.else_branch.unwrap().1 {
                            syn::Expr::Block(expr_block) => expr_block.block,
                            _ => unreachable!(),
                        };
                        visit_block(else_block, &mut false_block, mg);
                        let true_block = Rc::new(true_block);
                        let false_block = Rc::new(false_block);
                        let switch_int = SwitchInt {
                            rs: reg,
                            true_res: true,
                            true_block: Rc::clone(&true_block),
                            false_block: Rc::clone(&false_block),
                        };
                        this_block.set_end(BlockEnd::SwitchInt(switch_int));
                        mg.push_block(true_block);
                        mg.push_block(false_block);
                    }
                    syn::Expr::Return(return_expr) => {
                        match return_expr.expr {
                            Some(expr) => {
                                match visit_expr(*expr, this_block, mg) {
                                    Opt::Reg(reg) => {
                                        let inst = Inst {
                                            inst: InstOption::Ass,
                                            rd: Reg(0),
                                            rs1: reg.into(),
                                            rs2: reg.into(),
                                        };
                                        this_block.push_inst(inst);
                                    }
                                    Opt::Imm(imm) => {
                                        let inst = Inst {
                                            inst: InstOption::Ass,
                                            rd: Reg(0),
                                            rs1: imm.into(),
                                            rs2: imm.into(),
                                        };
                                        this_block.push_inst(inst);
                                    }
                                }
                                this_block.set_end(BlockEnd::Return);
                            }
                            None => {
                                this_block.set_end(BlockEnd::Return);
                            }
                        }
                    }
                    _ => unreachable!(),
                }
            }
            _ => unreachable!(),
        }
    }
    mg.tm.scope_pop();
}

fn visit_expr(expr: syn::Expr, block: &mut BasicBlock, mg: &mut MIRGen) -> Opt {
    match expr {
        syn::Expr::Binary(exp_bin) => {
            let l_opt = visit_expr(*exp_bin.left, block, mg);
            let r_opt = visit_expr(*exp_bin.right, block, mg);
            let reg = mg.create_reg();
            let inst_op = match exp_bin.op {
                syn::BinOp::Add(_) => InstOption::Add,
                syn::BinOp::Sub(_) => InstOption::Sub,
                syn::BinOp::Mul(_) => InstOption::Mul,
                _ => unreachable!(),
            };
            let inst = Inst {
                inst: inst_op,
                rd: reg,
                rs1: l_opt,
                rs2: r_opt,
            };
            block.push_inst(inst);
            reg.into()
        }
        syn::Expr::Lit(lit) => {
            match lit.lit {
                syn::Lit::Int(num) => {
                    num.base10_digits().parse::<i64>().unwrap().into()
                }
                _ => unreachable!(),
            }
        }
        syn::Expr::Path(path) => {
            // get ident
            let ident = path.path.get_ident().unwrap().to_string();
            mg.tm.sympol_get(&ident).expect(&format!("{} can't find", ident)).into()
        }
        _ => unreachable!(),
    }
}

fn visit_cond_expr(expr: syn::Expr, block: &mut BasicBlock, mg: &mut MIRGen) -> Reg {
    match expr {
        syn::Expr::Binary(exp_bin) => {
            let reg = mg.create_reg();
            let l_opt = visit_expr(*exp_bin.left, block, mg);
            let r_opt = visit_expr(*exp_bin.right, block, mg);
            let inst_op = match exp_bin.op {
                syn::BinOp::Gt(_) => InstOption::Gt,
                syn::BinOp::Lt(_) => InstOption::Lt,
                syn::BinOp::Ge(_) => InstOption::Ge,
                syn::BinOp::Le(_) => InstOption::Le,
                syn::BinOp::Eq(_) => InstOption::Eq,
                _ => unreachable!(),
            };
            let inst = Inst {
                inst: inst_op,
                rd: reg,
                rs1: l_opt,
                rs2: r_opt,
            };
            block.push_inst(inst);
            reg.into()
        }
        _ => unreachable!(),
    }
}
